package za.ac.sun.cs.viz.distrib;

import static org.abego.treelayout.internal.util.Contract.checkArg;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.abego.treelayout.util.AbstractTreeForTreeLayout;

public class VizTreeLayout<TreeNode> extends
		AbstractTreeForTreeLayout<TreeNode> {

	private Map<TreeNode, List<TreeNode>> childrenMap = new HashMap<TreeNode, List<TreeNode>>();

	private Map<TreeNode, TreeNode> parents = new HashMap<TreeNode, TreeNode>();

	public VizTreeLayout(TreeNode root) {
		super(root);
	}

	public void clear() {
		childrenMap = new HashMap<TreeNode, List<TreeNode>>();
		parents = new HashMap<TreeNode, TreeNode>();
	}

	@Override
	public List<TreeNode> getChildrenList(TreeNode node) {
		List<TreeNode> result = childrenMap.get(node);
		return result == null ? getEmptyList() : result;
	}

	@Override
	public TreeNode getParent(TreeNode node) {
		return parents.get(node);
	}

	private List<TreeNode> emptyList;

	private List<TreeNode> getEmptyList() {
		if (emptyList == null) {
			emptyList = new ArrayList<TreeNode>();
		}
		return emptyList;
	}

	public void addChild(TreeNode parentNode, TreeNode node) {
		checkArg(hasNode(parentNode), "parentNode is not in the tree");
		checkArg(!hasNode(node), "node is already in the tree");
		List<TreeNode> list = childrenMap.get(parentNode);
		if (list == null) {
			list = new ArrayList<TreeNode>();
			childrenMap.put(parentNode, list);
		}
		list.add(node);
		parents.put(node, parentNode);
	}

	private boolean hasNode(TreeNode node) {
		return node == getRoot() || parents.containsKey(node);
	}

}
